#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the version 3 of the GNU Lesser General Public License
#   as published by the Free Software Foundation.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU Lesser General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# Copyright (c) NEC Deutschland GmbH, NEC HPC Europe
#
# $Id$

from time import sleep, time
from json import dumps as json_dumps, loads as json_loads 
import pika
from pika.adapters import *
from pika.reconnection_strategies import SimpleReconnectionStrategy
from Queue import *
from threading import Thread
from traceback import *

from basic_types.marshaller import PrimitiveTypesMarshaller
from basic_types.marshaller import PrimitiveTypesUnmarshaller
from aggmon.channel import Channel
from aggmon.stoppable_thread import StoppableThread
from basic_types.factory import Factory


__all__ = ["PIKAChannel"]


#pika.log.setup(pika.log.DEBUG, color=True)
pika.log.setup(pika.log.INFO, color=True)


class PIKACallback(object):
    def __init__( self, notify_callback, workq, raw=False ):
        self.__workq = workq
        self.__notify_callback = notify_callback
        self.__raw = raw

    def callback( self, channel, method, properties, body):
        pika.log.info("Basic.Deliver %s delivery-tag %i: %s",
                  properties.content_type,
                  method.delivery_tag,
                  body)
        #print_stack()
        channel.basic_ack( delivery_tag=method.delivery_tag )

        # ignore special message that helps terminating the thread
        if body == "quit":
            # could do something smarter here?
            return
        if not self.__raw:
            jsonrpc_request = json_loads( body )
            message = Factory.build( jsonrpc_request["method"] )
            message.unmarshal( PrimitiveTypesUnmarshaller( jsonrpc_request["params"] ) )
        else:
            message = body
        topic = method.routing_key
        #print "unmarshalled message:", message
        self.__workq.put( {"cb": self.__notify_callback, "args": [message], "kwds": {"topic": topic} } )
        #self.__notify_callback( message, topic=topic )


class PIKAChannelExit(Exception):
    pass


class PIKAGenericChannel(StoppableThread):
    def __init__( self, connection=None, credentials=None, host=None, exchange_name=None, durable=False ):
        StoppableThread.__init__( self )
        self.durable = durable
        self.exchange_name = exchange_name
        self.connection = None
        self.channel = None
        self.__connecting = False
        self.__connected = False
        self.__credentials = credentials
        self.__host = host
        self._exchange_declared = False
        self._start_time = time()
        self.daemon = True
        if connection is None:
            self.connect()
        else:
            self.connection = connection
            self.on_connected( connection )

    def connect( self ):
        if self.__connecting:
            pika.log.info('PIKAChannel.connect: Already connecting to broker')
            return
        self.__connecting = True
        parameters = pika.ConnectionParameters( host=self.__host, port=5672, virtual_host="/", credentials=self.__credentials )
        self.connection = SelectConnection( parameters, self.on_connected )
        #self.connection = SelectConnection( parameters, self.on_connected, reconnection_strategy=SimpleReconnectionStrategy() )
        self.connection.add_on_close_callback( self.on_closed )
    
    def on_connected( self, connection ):
        pika.log.info("on_connected: Connected to broker on %s" % self.__host)
        self.connection = connection
        self.__connected = True
        connection.channel( self.on_channel_open )

    def on_channel_open( self, channel_ ):
        self.channel = channel_
        pika.log.info("on_channel_open: Received our channel")
        self.channel.exchange_declare( exchange=self.exchange_name, type="topic", auto_delete=False,
                                       durable=self.durable, callback=self.on_exchange_declared )

    def on_exchange_declared( self, frame ):
        pika.log.info("on_exchange_declared: got exchange")
        self._exchange_declared = True
        self.__connecting = False
        

    def on_basic_cancel( self, frame ):
        pika.log.info('PikaClient: Basic Cancel Ok')
        # If we don't have any more consumer processes running close
        self.connection.close()

    def on_closed( self, connection ):
        # We've closed our pika connection so stop the demo
        pika.log.info('on_closed: Basic Cancel Ok')
        self.__connected = False
        self.__connecting = False

    def _stopper( self ):
        if self._stopevent.isSet():
            print "Raising PIKAChannel Exit exception in", type(self), self.name
            raise PIKAChannelExit
        self.connection.ioloop.add_timeout( time() + 1, self._stopper )


    # Thread
    def run( self ):
        #sleep(1)
        self.connection.ioloop.add_timeout( time() + 1, self._stopper )
        while not self._stopevent.isSet():
            #print "in pika channel run loop", self._stopevent.isSet(), type(self)
            try:
                self.connection.ioloop.start()
            except PIKAChannelExit:
                print "Channel Exit!", type(self), self.name
                break
            except Exception, e:
                print "Exception in run loop"
                print e
                raise
        self.channel.close()
        print "Channel terminating...", type( self ).__name__, self.name
        #StoppableThread.join( self )


class PIKAWriteChannel(PIKAGenericChannel):
    def __init__( self, credentials=None, host=None, exchange_name=None, durable=False ):
        PIKAGenericChannel.__init__( self, credentials=credentials, host=host, exchange_name=exchange_name, durable=durable )
        self.pending = Queue()
        #self._timeout = self.connection.ioloop.add_timeout( 1000, self.send_pending )
        self.__timeout = None
        self.start()

    def cleanup( self ):
        # stop subscriber threads
        print "Cleanup called for write channel", type(self)
        self._stopevent.set()
        self.connection.ioloop.remove_timeout( self.__timeout )
        print "Cleanup finished for write channel"

    def publish( self, topic, message, application_headers={}, content_type="application/timacs-metric", raw=False, **kwds ):
        """
        Publish a message to the AMQP exchange the channel is pointing at.
        When the raw flag is True: interpret the message as payload. Otherwise encapsulate it into a JSON
        request, passing 'method', 'params' and 'id'.
        AMQP message properties such as content_type, app_id, type, reply_to, etc... can be passed along
        with the publish() call. They will be inserted into the AMQP message.
        """
        if raw:
            payload = message
        else:
            params = {}
            message.marshal( PrimitiveTypesMarshaller( params ) )
            jsonrpc_request = { "method": message.__class__.__name__.rsplit( "." )[-1], "params": params, "id": None }
            payload = json_dumps( jsonrpc_request )
            #print "publishing:", payload
        #properties = pika.BasicProperties( application_headers=application_headers, content_type=content_type, delivery_mode=1 )
        properties = pika.BasicProperties( content_type=content_type, delivery_mode=1 )
        msg = { "properties": properties,
                "body": str(payload),
                "routing_key": str(topic) }
        self.pending.put( msg )
        if self.__timeout is None:
            self.__timeout = self.connection.ioloop.add_timeout( 1, self.send_pending )

    def send_pending( self ):
        if self._exchange_declared and time() - self._start_time > 5:
            while self.pending.qsize() > 0:
                try:
                    msg = self.pending.get( True, 1 )
                except Empty:
                    print "Empty msg from pending"
                    break
                try:
                    #print "sending:", msg
                    self.channel.basic_publish( exchange=self.exchange_name, **msg )
                except Exception, e:
                    print "Exception in publish!"
                    print e
                    raise
                self.pending.task_done()
        self.__timeout = self.connection.ioloop.add_timeout( 1, self.send_pending )

        
class PIKASubscriptionChannel(object):
    def __init__( self, connection, channel, exchange_name, callback=None, durable=False, topic="#" ):
        self.callback = callback
        self.connection = connection
        self.channel = channel
        self.exchange_name = exchange_name
        self.topic = topic
        self.channel.queue_declare( durable=durable, exclusive=True, auto_delete=True,
                                    callback=self.on_queue_declared )

    def cleanup( self ):
        if self.queue_name is not None:
            self.channel.basic_cancel( self.queue_name )
            #self.channel.queue_delete( self.queue_name )

    def on_queue_declared(self, frame):
        self.queue_name = frame.method.queue
        pika.log.info( "on_queue_declared: Queue %s declared, Binding Queue ... %s" % (self.queue_name, str(frame)) )
        self.channel.queue_bind( exchange=self.exchange_name,
                                 queue=self.queue_name,
                                 routing_key=self.topic,
                                 callback=self.on_queue_bound )

    def on_queue_bound( self, frame ):
        pika.log.info('on_queue_bound: Queue bound for subscriber, Issuing Basic Consume')
        self.channel.basic_consume( consumer_callback=self.callback,
                                    queue=self.queue_name, no_ack=False )


class PIKAReadChannel(PIKAGenericChannel):
    def __init__( self, connection=None, credentials=None, host=None, exchange_name=None, durable=False, workq=None ):
        PIKAGenericChannel.__init__( self, connection=connection, credentials=credentials, host=host, exchange_name=exchange_name, durable=durable )
        self.__issued_subscribers = Queue()
        self.__subscribers = list()
        self.__workq = workq
        self.__timeout = None
        self.start()

    def cleanup( self ):
        # stop subscriber threads
        print "Cleanup called for read channel", type(self)
        self.connection.ioloop.remove_timeout( self.__timeout )
        for subscriber in self.__subscribers:
            subscriber.cleanup()
        self._stopevent.set()
        print "Cleanup finished for read channel"

    def handle_subscriptions( self ):
        if not self._exchange_declared or self.__issued_subscribers.qsize() == 0:
            self.__timeout = self.connection.ioloop.add_timeout( time() + 1, self.handle_subscriptions )
            return
        try:
            sub = self.__issued_subscribers.get()
        except Empty:
            print "no subscriber"
        else:
            self.__subscribers.append( PIKASubscriptionChannel( self.connection, self.channel, self.exchange_name, **sub ) )
        self.__timeout = self.connection.ioloop.add_timeout( self.__issued_subscribers.qsize() > 0 and 1 or time() + 1,
                                                             self.handle_subscriptions )

    def subscribe( self, topic, notify_callback, raw=False ):
        pika.log.info( "subscribing " + str(type(notify_callback)) + " with topic " + topic )
        # put subscriber onto wait list
        cb = PIKACallback( notify_callback, self.__workq, raw=raw )
        self.__issued_subscribers.put( {"topic": topic, "callback": cb.callback } )
        if self.__timeout is None:
            self.__timeout = self.connection.ioloop.add_timeout( 1, self.handle_subscriptions )

class PIKAChannel(Channel):
    def __init__( self, url, durable=False ):
        Channel.__init__( self )
        self.url = url
        self.durable = durable

        scheme, rest = url.split( "://" )
        assert scheme == "pika"
        netloc, channel_name = rest.split( "/" )
        userid_password, self.__host = netloc.split( "@" )
        userid, password = userid_password.split( ":" )
        #self.__channel_name = "/data/" + channel_name 
        self.exchange_name = channel_name
        credentials = pika.PlainCredentials( userid, password )

        self.callback_work = Queue()
        # start exactly one worker thread to not break the serialization of the database
        # We want a worker thread here because the foreign callbacks should not be done in
        # the path of the connection's ioloop.
        self.cb_thread = Thread( target=self.cb_worker, name="PikaWrk_" + self.url )
        self.cb_thread.daemon = True
        self.cb_stopping = False
        self.cb_thread.start()

        self.w_channel = PIKAWriteChannel( credentials=credentials, host=self.__host,
                                           exchange_name=self.exchange_name, durable=durable )
        while not self.w_channel._exchange_declared:
            continue
        self.r_channel = PIKAReadChannel( connection=None, credentials=credentials, host=self.__host,
                                          exchange_name=self.exchange_name, durable=durable,
                                          workq=self.callback_work )
        #self.w_channel = self.r_channel
        # statistics
        self.cnt_publish = 0
        self.cnt_subscribers = 0
        self.called_subscribers = 0
        self.called_callbacks = {}

    def cb_worker( self ):
        pika.log.info( "started callback worker thread for channel %s" % self.url )
        while not self.cb_stopping:
            work = self.callback_work.get()
            function = work["cb"]
            args = work["args"] or []
            kwds = work["kwds"] or {}
            self.called_subscribers += 1
            cb_name = function.im_class.__name__.split(".")[-1] + "_" + function.__name__
            if cb_name in self.called_callbacks:
                self.called_callbacks[cb_name] += 1
            else:
                self.called_callbacks[cb_name] = 1
            pika.log.info( "****** WORKER THREAD has work!!! Calling %s" % cb_name )
            function( *args, **kwds )
            self.callback_work.task_done()
        pika.log.info( "stopping callback worker thread for channel %s" % self.url )
        self.cb__thread.join()

    def cleanup( self ):
        print "Cleanup called for channel", self.url
        self.r_channel.cleanup()
        self.w_channel.cleanup()
        self.cb_stopping = True

    def publish( self, topic, message, application_headers={}, content_type="application/timacs-metric", raw=False, **kwds ):
        self.cnt_publish += 1
        self.w_channel.publish( topic, message, application_headers=application_headers,
                                content_type=content_type, raw=raw, **kwds )

    def stats( self ):
        return { "url": self.url,
                 "subscribers": self.cnt_subscribers,
                 "published": self.cnt_publish,
                 "called_subscribers": self.called_subscribers,
                 "callbacks": self.called_callbacks,
                 "queue_length": self.callback_work.qsize() }

    def subscribe( self, topic, notify_callback, raw=False ):
        self.cnt_subscribers += 1
        self.r_channel.subscribe( topic, notify_callback, raw=raw )

